#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@创建时间    : 2023/3/9  13:58
@作者  : st
@文件名: entity_replenish.py
@项目名: PyCharm
@文件描述:
    针对个别识别错误的实体进行补全
"""

import inspect
import os

abspath = os.path.abspath(inspect.getfile(inspect.currentframe()))
root_path = os.path.dirname(os.path.dirname(abspath))


class EntityReplenish:
    def __init__(self):

        temp_dict = dict()
        for line in open(os.path.join(root_path, 'data', 'dict', '缺失实体映射表.txt'), 'r', encoding='utf-8'):
            word, word_whole = line.strip().split('\t')
            if not temp_dict.__contains__(word):
                temp_dict[word] = set()
            temp_dict[word].add(word_whole)
        self.entity_repl_dict = dict()
        for key in temp_dict.keys():
             self.entity_repl_dict[key] = sorted(temp_dict[key], key=lambda x: len(x), reverse=True)
        print('loaded 缺失实体映射表')
        self.entity_dict = dict()
        for line in open(os.path.join(root_path, 'data', 'dict', '实体表.txt'), 'r', encoding='utf-8'):
            word, label = line.strip().split('\t')
            if not self.entity_dict.__contains__(label):
                self.entity_dict[label] = set()
            self.entity_dict[label].add(word)
        for label in self.entity_dict.keys():
            temp_list = list(self.entity_dict[label])
            self.entity_dict[label] = sorted(temp_list, key=lambda x: len(x), reverse=True)
        print('loaded 实体表')

    def entity_full_batch(self, entities, text):
        """
        实体填充
        """
        for label in entities.keys():
            entitys = entities[label]
            new_entity_list = []
            for entity in entitys:
                if self.entity_repl_dict.__contains__(entity):
                    for word in self.entity_repl_dict[entity]:
                        if word in text:
                            entity = word
                            break
                if not new_entity_list.__contains__(entity):
                    new_entity_list.append(entity)
            entities[label] = new_entity_list
        for key in self.entity_dict.keys():
            for word in self.entity_dict[key]:
                if word in text:
                    if not entities.__contains__(key):
                        entities[key] = [word]
                    elif not entities[key].__contains__(word):
                        entities[key].append(word)
        return entities

